%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%
% This file uses the Case II version of the Bansal and Yaron (2004)
% model (with volatility risk). The default frequency is monthly.
%
% Global settings:
T = 1000; % number of simulations
L = 120000; % maximum length of recursive utility (10,000 years)
N = 100; % number of iteration for the vc formula
H = 600;
Draw = 1000; % number of draws to obtain expected utility
rng(120523); % random seed (May 12, 2023)
Progress = 1; % progress bar 'on' = 1, otherwise 0; TPon = 1 required
%
% Required packages:
% Parallel Computing Toolbox
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Part 0: Input Parameters
% Assumed coefficients
mu = 0.0015; % mean consumption and dividend growth %
sigmabar = 0.0072; % baseline standard deviation of consumption growth innovation % sigmabar = 0.0078
phic = 1; % consumption growth multiplier of the latent process
phid = 2.5; % dividend growth multiplier of the latent process % phi = 3
alphax = 0.038; % volatility impact factor of the latent process % phie = 0.045
alphac = 1; % volatility impact factor of the consumption growth innovation
alphad = 5.96; % volatility impact factor of the dividend growth innovation % phid = 4.5
gamma = 11; % coefficient of relative risk aversion % gamma = 11/10/7.5
gammatilde = 3; % coefficient of relative risk aversion in future periods % gammatilde = 5.3
rho = 1/2; % intertemporal elasticity of substitution % psi = 1.5
nux = 0.975; % persistence coefficient of the latent process % rho = 0.979
beta = 0.9989; % subjective discount factor % delta = 0.998
nusigma = 0.999; % persistence coefficient of volatility % nu1 = 0.987
alphasigma = 0.0000028; % standard deviation of volatility % sigmaw = 0.0000023
chi = 2.6; % dividend consumption exposure % pi = 0

% Implied settings
Y = floor(H/12); % maximum horizon in years

%% Part 1: Preliminary Processing
% Initialisation
sigmasq0 = sigmabar^2; % conditional volatility of consumption
x0 = 0; % varying component of the expected consumption growth rate 
c0 = 0; % initial log consumption (without loss of generality)
d0 = 0; % initial log dividend (without loss of generality)

% Empty vectors
sigmasq = zeros(T,1); % conditional volatility of consumption
x = zeros(T,1); % varying component of the expected consumption growth rate 
c = zeros(T,1); % log consumption
d = zeros(T,1); % log dividend

%% Part 2: Fixed Point Estimation
% Coefficients of vc
A1v = (beta*phic)/(1-beta*nux); % as defined in AES2023 Appendix, phi_v
A2v = 0.5*beta*(1-gammatilde)*((alphac^2)+(A1v^2)*(alphax^2))/(1-beta*nusigma); % psi_tilde_v
A2v_2 = 0.5*beta*((alphac^2)+(A1v^2)*(alphax^2))/(1-beta*nusigma); % psi_tilde_v/(1-gammatilde)
A0v = (beta/(1-beta))*(mu+A2v_2*(1-gammatilde)*(sigmabar^2)*(1-nusigma)+0.5*(1-gammatilde)^3*(A2v_2^2)*(alphasigma^2)); % mu_tilde_v

% Capital market equilibrium (AES version)
fun = @(z) AES2019P2(z,alphac,alphad,alphax,alphasigma,gamma,gammatilde,rho,mu,nux,nusigma,sigmabar,phic,phid,chi,A1v,A2v_2);
zbar = fsolve(fun,0); % equilibrium pd ratio

% Coeffieients at equilibrium point
pibar = -mu-0.5*(1-rho)*((1-gamma)+nusigma*(gamma-gammatilde))*(alphac^2+A1v^2*alphax^2)*sigmabar^2 ...
    -0.5*(1-gamma)^2*A2v_2^2*(1-gammatilde)^2*alphasigma^2; % equilibrium log SDF
kappa1 = exp(zbar)/(1+exp(zbar)); % kappa from CS1988
kappa0 = log(1+exp(zbar))-kappa1*zbar;
A1 = (phid-rho*phic)/(1-kappa1*nux); % as defined in AES2023 appendix
A2 = (-0.5*((rho-gamma)*(1-gamma)-(1-rho)*(gamma-gammatilde)*nusigma)*(alphac^2+A1v^2*alphax^2)...
    +0.5*(chi-gamma)^2*alphac^2+0.5*(kappa1*A1+(rho-gamma)*A1v)^2*alphax^2 ...
    +0.5*alphad^2)/(1-kappa1*nusigma);
A0 = (1/(1-kappa1))*(pibar+kappa0+kappa1*A2*sigmabar^2*(1-nusigma) ...
    +0.5*(((rho-gamma)*(1-gammatilde)+(1-rho)*(gamma-gammatilde)*(1-nusigma))*A2v_2+kappa1*A2)^2*alphasigma^2);

%% Part 3: Simulation
% Random shocks
epsilon = randn(T,4); % 4 columns correspond to w (sigma), e (x), u (d), and eta (c), respectively

% Volatility
sigmasq(1) = sigmabar^2+nusigma*(sigmasq0-sigmabar^2)+alphasigma*epsilon(1,1); % first value (period t+1)
for t=2:T
    sigmasq(t) = sigmabar^2+nusigma*(sigmasq(t-1)-sigmabar^2)+alphasigma*epsilon(t,1);
    if sigmasq(t)<sigmabar^2/10000
        sigmasq(t) = sigmabar^2/10000;
    end
end

% Varying component of the expected consumption growth rate 
x(1) = nux*x0+alphax*(sigmasq0^0.5)*epsilon(1,2); % first value (period t+1)
for t=2:T
    x(t) = nux*x(t-1)+alphax*(sigmasq(t-1)^0.5)*epsilon(t,2);
end

% Consumption
c(1) = c0+mu+phic*x0+alphac*(sigmasq0^0.5)*epsilon(1,4); % first value (period t+1)
for t=2:T
    c(t) = c(t-1)+mu+phic*x(t-1)+alphac*(sigmasq(t-1)^0.5)*epsilon(t,4);
end

% Dividend
d(1) = d0+mu+phid*x0+alphad*(sigmasq0^0.5)*epsilon(1,3)+chi*alphac*(sigmasq0^0.5)*epsilon(1,4);
for t=2:T
    d(t) = d(t-1)+mu+phid*x(t-1)+chi*alphac*(sigmasq(t-1)^0.5)*epsilon(t,3)+alphad*(sigmasq(t-1)^0.5)*epsilon(t,4);
end

% Ratios
vc = A0v+A1v*x+A2v_2*(1-gammatilde)*sigmasq; % log value-consumption ratio
z = ones(T,1)*A0+A1*x+A2*sigmasq; % log price dividend ratio

% Intermediary parameters
A = alphac^2+(A1v^2)*(alphax^2);
B = -0.5*((rho-gamma)*(1-gamma)-(1-rho)*(gamma-gammatilde)*nusigma)*A;

% Initialisation 
psi_b_tilde_v = zeros(H,1); % \tilde\psi_{b,h}
psi_b_v = zeros(H,1); % psi_{b,h}
psi_d_tilde_v = zeros(H,1); % \tilde\psi_{d,h}
phi_b_v = zeros(H,1); % \phi_{b,h}
phi_d_v = zeros(H,1); % \phi_{d,h}
psi_v = zeros(H,1); % \psi_h
Psi_v = zeros(H,1); % \Psi_h
phi_v = zeros(H,1); % \phi_h
Phi_v = zeros(H,1); % \Phi_h
mu_d_tilde_v = zeros(H,1); % \mu_d,h
mu_b_tilde_v = zeros(H,1); % \mu_b,h
nux_v = zeros(H,1); % (1-nv_x^h)/(1-nv_x)
a_v = zeros(H,1); % a_h
b_v = zeros(H,1); % b_h

% Parameters of different horizons
psinu = -0.5*(gamma-1)*A/(1-nusigma);
% psinutilde  = -0.5*(gammatilde-1)*A/(1-nusigma);
for h=1:H
    nux_v(h) = (1-nux^(h))/(1-nux);
end
for h=1:H
    phi_b_v(h) = -rho*phic*nux_v(h);
    phi_d_v(h) = (-rho*phic+phid)*nux_v(h);
    phi_v(h) = gamma*chi*(alphac^2) + (gamma-rho*(nux^(h-1)))*(1-nux^(h-1))*phic*phid*(alphax^2)/((1-nux)^2);
    Phi_v(h) = (chi^2)*(alphac^2) + (nux_v(h)^2)*phid^2*alphax^2 + alphad^2;
end
psi_b_tilde_v(1) = B + 0.5*((gamma^2)*(alphac^2)+((rho-gamma)*A1v)^2*alphax^2);
psi_b_v(1) = psi_b_tilde_v(1) - 0.5*(1-rho)*(gamma-gammatilde)*nusigma*A;
psi_d_tilde_v(1) = B + 0.5*(((chi-gamma)*alphac)^2+((rho-gamma)*A1v*alphax)^2+alphad^2);
for h=2:H
    psi_b_tilde_v(h) = B + psi_b_tilde_v(h-1)*nusigma + 0.5*((gamma^2)*(alphac^2)+((rho-gamma)*A1v+phi_b_v(h-1))^2*alphax^2);
    psi_b_v(h) = psi_b_tilde_v(h) - 0.5*(1-rho)*(gamma-gammatilde)*nusigma*A*nux_v(h);
    psi_d_tilde_v(h) = B + psi_d_tilde_v(h-1)*nusigma + 0.5*(((chi-gamma)*alphac)^2+(((rho-gamma)*A1v+phi_d_v(h-1))*alphax)^2+alphad^2);
end
Psi_v(1) = chi^2*alphac^2;
mu_d_tilde_v(1) = -0.5*(1-rho)*A*((1-gamma)+nusigma*(gamma-gammatilde))*sigmabar^2 + 0.5*((((rho-gamma)*(1-gammatilde)+(1-rho)*(gamma-gammatilde)*(1-nusigma))*A2v_2)^2-(1-gamma)^2*A2v_2^2*(1-gammatilde)^2)*alphasigma^2;
mu_b_tilde_v(1) = - mu + (-0.5*(1-rho)*((1-gamma)+nusigma*(gamma-gammatilde))*(alphac^2+A1v^2*alphax^2))*sigmabar^2 + 0.5*((((rho-gamma)*(1-gammatilde)+(1-rho)*(gamma-gammatilde)*(1-nusigma))*A2v_2)^2-(1-gamma)^2*(A2v_2*(1-gammatilde))^2)*alphasigma^2;
for h=2:H
    psi_v(h) = psi_b_tilde_v(h-1) - psi_d_tilde_v(h-1);
    Psi_v(h) = chi^2*alphac^2 + (phi_d_v(h-1)-phi_b_v(h-1))^2*alphax^2;
    mu_d_tilde_v(h) = mu_d_tilde_v(h-1) + (psi_d_tilde_v(h-1)*(1-nusigma) ...
        -0.5*(1-rho)*A*((1-gamma)+nusigma*(gamma-gammatilde)))*sigmabar^2 ...
        +0.5*((((rho-gamma)*(1-gammatilde)+(1-rho)*(gamma-gammatilde)*(1-nusigma))*A2v_2+psi_d_tilde_v(h-1))^2-(1-gamma)^2*A2v_2^2*(1-gammatilde)^2)*alphasigma^2;
    mu_b_tilde_v(h) = mu_b_tilde_v(h-1) - mu + (psi_b_tilde_v(h-1)*(1-nusigma)-0.5*(1-rho)*((1-gamma)+nusigma*(gamma-gammatilde))*(alphac^2+A1v^2*alphax^2))*sigmabar^2 ...
        +0.5*((((rho-gamma)*(1-gammatilde)+(1-rho)*(gamma-gammatilde)*(1-nusigma))*A2v_2+psi_b_tilde_v(h-1))^2-(1-gamma)^2*(A2v_2*(1-gammatilde))^2)*alphasigma^2;
end
for h=1:H
    a_v(h) = (gamma-rho)*abs(psinu) + psi_b_v(h);
    b_v(h) = 0.5*A*((gamma-1)+(1-rho)*(nusigma^h))/(1-nusigma);
end

% Clear innitial variables
clear fun sigmasq0 c0 d0 x0 cov_rz

% Clear intermediate variables
clear t A B B_0 cons cons1 cons1_0 EP_uncon_0 ler sr rbar_0 rbar_fsigma

%% Part 4: Timing Premium
% Initial values
rho2 = [rho 1]';
tp = zeros(T,size(rho2,1));
TP = zeros(size(rho2));

% Scenario 1: rho=1
tp(:,end) = repmat(beta/(1-beta)*A2v_2*sigmabar^2*(1-nusigma)*...
    ((1-gammatilde)-(1-gamma)*(1-beta*nusigma)*beta/((1-beta^2*nusigma)*(1+beta))),T,1)+...
    A2v_2*nusigma*sigmasq*beta*((1-gamma)*(1-beta)/(1-beta^2*nusigma)+gamma-gammatilde)+...
    repmat(0.5*beta/(1-beta)*(1-gamma+beta*(gamma-gammatilde))*A2v_2^2*(1-gammatilde)^2*alphasigma^2,T,1); % log time premium
TP(end) = mean(ones(T,1)-exp(tp(:,end))); % time premium 

% Scenario 2: rho=\1 (estimating vc function)
I = size(rho2,1)-1; % number of index
parfor h=1:I
    % Initialisation
    gap_x = (alphax*sigmabar)/(1-nux^2)^0.5/6; % adjustment coefficient 6
    Grid_x = gap_x*(-19.5:1:19.5)';
    gap_sigmasq = alphasigma/(1-nusigma^2)^0.5/14; % adjustment coefficient 14
    Grid_sigmasq = sigmabar^2*ones(1,40)+gap_sigmasq*(-11:28);

    % vc as a function of x and sigmasq (rho=1)
    vc0 = repmat(A0v,40,40)+repmat(A1v*Grid_x,1,40) ...
        +repmat(A2v_2*(1-gammatilde)*(Grid_sigmasq),40,1);
    vc0 = [vc0 vc0(:,end)]; % replicate column 41
    vc0 = [vc0;vc0(end,:)]; % replicate row 41
    
    % vc as a function of x and sigmasq (rho\=1)
    for n=1:N % T times of iteration
        for i=1:40 % all grid points
            for j=1:40

                % Second period
                epsilon_d = randn(Draw,3);
                vc1 = zeros(Draw,1); % next period vc
                x_d = nux*Grid_x(i)+alphax*(Grid_sigmasq(j)^0.5)*epsilon_d(:,1);
                sigmasq_d = sigmabar^2+nusigma*(Grid_sigmasq(j)-sigmabar^2)+alphasigma*epsilon_d(:,2);
                sigmasq_d = max([repmat(sigmabar^2/10000,Draw,1) sigmasq_d],[],2);
                for draw=1:Draw
                    f_x = min(floor((x_d(draw)+20.5*gap_x)/gap_x),40); % position, max 40                 
                    f_sigmasq = min(floor((sigmasq_d(draw)-sigmabar^2+12*gap_sigmasq)/gap_sigmasq),40); % max 40                   
                    if f_x>0 && f_sigmasq>0
                        d_x = (x_d(draw)-Grid_x(f_x,1))/gap_x; % distance from the grid (0 to 1)
                        d_sigmasq = (sigmasq_d(draw)-Grid_sigmasq(f_sigmasq))/gap_sigmasq;
                        vc1(draw) = vc0(f_x,f_sigmasq)+(vc0(f_x+1,f_sigmasq)-vc0(f_x,f_sigmasq))*(1-d_x)+(vc0(f_x+1,f_sigmasq+1)-vc0(f_x+1,f_sigmasq))*(1-d_sigmasq);   
                    elseif f_x<=0 && f_sigmasq>0
                        d_sigmasq = (sigmasq_d(draw)-Grid_sigmasq(f_sigmasq))/gap_sigmasq;
                        vc1(draw) = vc0(1,f_sigmasq)+(vc0(1,f_sigmasq+1)-vc0(1,f_sigmasq))*(1-d_sigmasq); 
                    elseif f_x>0 && f_sigmasq<=0
                        d_x = (x_d(draw)-Grid_x(f_x,1))/gap_x; % distance from the grid (0 to 1)
                        vc1(draw) = vc0(f_x,1)+(vc0(f_x+1,1)-vc0(f_x,1))*(1-d_x); 
                    else
                        vc1(draw) = vc0(1,1);
                    end
                end

                % First period
                vc0(i,j) = 1/(1-rho2(h))*log(1-beta+beta*(mean(exp((1-gammatilde)*vc1))* ...
                    mean(exp((mu+phic*Grid_x(i)+alphac*Grid_sigmasq(j)*epsilon_d(:,3))*(1-gammatilde))))^((1-rho2(h))/(1-gammatilde)));

                % monotonicity restriction 
                vc0(1:i-1,j) = min(vc0(1:i-1,j),vc0(i,j)); % vc increasing in x
                vc0(i+1:end,j) = max(vc0(i+1:end,j),vc0(i,j));
            end
        end

        % 41st row and coulumn
        vc0(end,1:end-1) = vc0(end-1,1:end-1);
        vc0(:,end) = vc0(:,end-1);
    end

    % Timing premium
    [TP(h),tp(:,h)] = AES2019P5(rho2(h),h,I,T,L,Draw,Progress,c,d,x,vc0,sigmasq,alphac,alphad,alphax,alphasigma,beta,gamma,mu,nux,nusigma,sigmabar,phic,phid,chi,gap_x,gap_sigmasq,Grid_x,Grid_sigmasq);
end

% Intermediary variables
clear I tp
